module TypeChecker where

import AbsJavalette
import PrintJavalette
import ErrM
import Control.Monad(foldM)
import qualified Data.Map as Map
import Data.Map (Map, empty, insert, fromList)


type Env = (Sigs, [Context])
type Sigs = Map Ident ([Type], Type)
type Context = Map Ident Type

emptyEnv :: Env
emptyEnv = (sigs, context) where
	sigs = fromList [(Ident "printInt", ([Int], Void)), (Ident "printString", ([String], Void)), (Ident "printDouble", ([Doub], Void)), (Ident "readInt", ([], Int)), (Ident "readDouble", ([], Doub))]
	context = []


lookupVar :: Env -> Ident -> Err Type
lookupVar (s, []) identifier = fail ("No such variable " ++ show identifier)
lookupVar (s, curr:rest) identifier = case Map.lookup identifier curr of
	Nothing -> lookupVar (s, rest) identifier
	Just t -> return t


--TODO add Err ([Type], Type)
--Takes an environment and an identifier of a function and returns 
--the functions signature if exists or an error
lookupFun :: Env -> Ident -> Err ([Type], Type)
lookupFun (signatures, _) identifier = case Map.lookup identifier signatures of
	Nothing -> fail ("function not defined " ++ show identifier)
	Just (args, retType) -> return (args, retType)


--TODO change to use err Env return type
--top is the top of the context stack, that is, the current scope
--TODO error if variable is redeclared, if exists in scope then error
updateVar :: Env -> Ident -> Type -> Err Env
updateVar (sigs, top:rest) identifier typ = case Map.lookup identifier top of
	Nothing -> return (sigs, (insert identifier typ top):rest)
	Just t -> fail "Redeclaration of variable in same scope"


collectTypes :: [Arg] -> [Type]
collectTypes [] = []
collectTypes ((Arg t _):rest) = t:(collectTypes rest)

updateFun :: Env -> TopDef -> Err Env
updateFun (sigs, context) (FnDef retType identifier args _) = 
	return (insert identifier (collectTypes args, retType) sigs, context)

--When entering a new block, a new empty context is pushed on the stack
newBlock:: Env -> Env
newBlock (sigs, context) = (sigs, empty:context)

checkExp :: Env -> Type -> Expr -> Err Expr
checkExp env typ expr = do
         typ2 <- inferExp env expr
         if (typ2 == typ) 
            then
                return expr
            else
                fail $ "Type of " ++ printTree expr ++
                       ", expected '" ++ printTree typ ++ "'" ++
                       " but found '" ++ printTree typ2 ++ "'."



checkStmt :: Type -> Env -> Stmt -> Err Env
checkStmt typ environment x = case x of
	Empty -> return environment
	BStmt (Block statements) -> case foldM (checkStmt typ) (newBlock environment) statements of
		Ok (sig, ctx_start:rest) -> return (sig, rest) --pop context of environment on exiting code block
		Bad s -> fail s
	Decl t items ->  foldM checkDecl environment items where
		checkDecl :: Env -> Item -> Err Env
		checkDecl environment (NoInit identifier) = updateVar environment identifier t
		checkDecl environment (Init identifier expression) = case checkExp environment t expression of
			Ok _ -> updateVar environment identifier t
			Bad s -> fail s
	Ass identifier expression -> case lookupVar environment identifier of
		Ok t -> case checkExp environment t expression of
			Ok _ -> return environment
			Bad s -> do
				fail s
		Bad s -> fail s
	Incr identifier -> case lookupVar environment identifier of
		Ok Int -> return environment
		--Ok t -> fail s
		Bad s -> fail s
		t -> fail "Can't increment on non integers"
	Decr identifier -> case lookupVar environment identifier of
		Ok Int -> return environment
		Bad s -> fail s
		t -> fail "Can't decrement on non integers"
	Ret expression -> case checkExp environment typ expression of
		Ok _ -> return environment
		Bad s -> fail s
	VRet -> if typ == Void
		then return environment
		else fail "Non void return"
	Cond expression statement -> case checkExp environment Bool expression of
		Ok _ -> checkStmt typ (newBlock environment) statement
		Bad s -> fail "Non boolean conditional expression in if statement"
	CondElse expression statement1 statement2 -> case checkExp environment Bool expression of
		Ok _ -> do
			checkStmt typ (newBlock environment) statement1
			checkStmt typ (newBlock environment) statement2
		Bad s -> fail "Non boolean conditional expression in if/else statement"
	While expression statement -> case checkExp environment Bool expression of
		Ok _ -> checkStmt typ (newBlock environment) statement
		Bad s -> fail "Non boolean conditional in while statement"
	SExp expression -> case inferExp environment expression of
		Ok _ -> return environment
		Bad s -> fail s



updateEnvWithArg :: Env -> Arg -> Err Env
updateEnvWithArg environment (Arg typ identifier) = case updateVar environment identifier typ of
	Ok e -> return e
	Bad s -> fail s

--Any return statements should be of the same type as the function
--Extend the environment with arguments and types
--
checkDef :: Env -> TopDef -> Err Env
checkDef environment (FnDef typ identifier args statements) = case foldM updateEnvWithArg (newBlock environment) args of
	Ok env' -> case  checkStmt typ env' (BStmt statements) of
		Ok _ -> return environment
		Bad s -> fail s
	Bad s -> fail s

--newBlock
--updateVar env, argname, argtyp
--checkStmt rettype env funcBlock

checkProgram :: Program -> Err Program
checkProgram (Program fdefs) = case foldM updateFun (emptyEnv) fdefs of
	Ok e -> case foldM checkDef e fdefs of
		Ok _ -> return (Program fdefs)
		Bad s -> fail s
	Bad s -> fail s

--inferExp :: Env -> Expr -> Err Type
inferExp :: Env -> Expr -> Err Type
inferExp env x = case x of
	ELitTrue -> return Bool
	ELitFalse -> return Bool
	ELitInt _ -> return Int
	ELitDoub _ -> return Doub
	EString _ -> return String
	EVar identifier -> case lookupVar env identifier of
		Ok t -> return t
		Bad s -> fail s
	Neg expression -> case inferExp env expression of
		Ok t -> return t
		Bad s -> fail s
	Not expression -> case inferExp env expression of
		Ok t -> return t
		Bad s -> fail s --maybe we should only return if expression is Bool
	EMul leftE op rightE -> case inferExp env leftE of --this might work, get type of left hand side, check if that is the same as the type of the right hand side expression
		Ok t -> case checkExp env t rightE of
			Ok _ -> return t
			Bad s -> fail s
		Bad s -> fail s
	EAdd leftE op rightE -> case inferExp env leftE of
		Ok t -> case checkExp env t rightE of
			Ok _ -> return t
			Bad s -> fail s
		Bad s -> fail s
	ERel leftE op rightE -> case inferExp env leftE of --types on both sides needs to be the same in order to compare, if so the return type is boolean
		Ok t -> case checkExp env t rightE of
			Ok _ -> return Bool
			Bad s -> fail s
		Bad s -> fail s
	EAnd leftE rightE -> case checkExp env Bool leftE of
		Ok e -> case checkExp env Bool rightE of
			Ok e -> return Bool
			Bad s -> fail s
		Bad s -> fail s
	EOr leftE rightE -> case checkExp env Bool leftE of
		Ok e -> case checkExp env Bool rightE of
			Ok e -> return Bool
			Bad s -> fail s
		Bad s -> fail s
	EApp identifier expressions -> case lookupFun env identifier of
		Ok sig -> return (snd sig)
		Bad s -> fail s



--bad008.jl passes, non reachable return
--bad015.jl passes, typecheck on parameter to function does not fail
--bad021.jl passes, void return in main function
--bad024.jl passes





--intended function: compile :: Stirng -> IO()
--compile :: Env -> Type -> Expr -> IO ()
--compile ::  Env -> Ident -> IO ()
--compile env identifier = case lookupFun env identifier of
--	Bad err -> do
--		putStrLn "Type error!"
--		putStrLn err
--	Ok _ -> do
--		putStrLn "Ok"

--version which is given in the book of last course
--compile :: String -> IO ()
--compile inputString = case checkProgram (myLexer inputString) of
--        Bad err -> do
--                putStrLn "SYNTAX ERROR"
--                putStrLn err
--        Ok tree -> case typecheck tree of
--           Bad err -> do
--                   putStrLn "TYPE ERROR"
--                   putStrLn err
--           Ok _ -> putStrLn "OK" --or go to next compiler phase

--testLookupVar :: Env -> Ident -> IO ()
--testLookupVar env identifier = case lookupVar env identifier of
--	Bad err -> do
--		putStrLn "No such var"
--		putStrLn err
--	Ok _ -> do
--		putStrLn "Ok"


--main = do
--	testLookupVar emptyEnv (Ident "i")


--Program [FnDef Int (Ident "main") [] (Block [Ret (ELitInt 0)])]


--checkProgram :: Program -> Err Program